// Copyright 2023 Srhino Co., Ltd. All rights reserved.

package v2

import (
	"context"
	"encoding/json"
	"errors"
	v2 "git.ouryun.cn/lzj1/grpc-bidirectional-stream/pkg/apis/v2"
	ifacev2 "git.ouryun.cn/lzj1/grpc-bidirectional-stream/pkg/interfaces/v2"
	"git.ouryun.cn/lzj1/grpc-bidirectional-stream/pkg/utils"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"io"
	"log"
	"net"
	"sync/atomic"
	"time"
)

var reqId int64

type grpcServer struct {
	v2.UnimplementedStreamDataServiceServer
	address    string
	streamData *StreamData
	streams    map[string]v2.StreamDataService_AgentStreamServer
	callback   RequestProcessFunc
	ctx        context.Context
	serverType string
}

type RequestProcessFunc func(message *v2.StreamMessage)

var _ ifacev2.GrpcServer = (*grpcServer)(nil)

func NewGrpcServer(addr string, serverType string, ctx context.Context, callback RequestProcessFunc) ifacev2.GrpcServer {
	return &grpcServer{
		address:    addr,
		streamData: NewStreamData(),
		callback:   callback,
		streams: map[string]v2.StreamDataService_AgentStreamServer{
			CLIENT_TYPE_WEBSERVER: nil,
			SERVER_TYPE_SC:        nil,
		},
		ctx:        ctx,
		serverType: serverType,
	}
}

func (s *grpcServer) AgentStream(stream v2.StreamDataService_AgentStreamServer) error {
	defer func() {
		if r := recover(); r != nil {
			log.Println(r)
		}
	}()
	utils.Logger.Info("client connection is successful")

	// 获取客户端的唯一标识符，可以根据实际需求进行处理

	md, ok := metadata.FromIncomingContext(stream.Context())
	marshal, _ := json.Marshal(md)
	utils.Logger.Infof("client id is %s", marshal)
	if !ok {
		return status.Errorf(codes.InvalidArgument, "missing metadata")
	}

	val := md.Get(CLIENT_KEY)
	if len(val) == 0 {
		return status.Errorf(codes.InvalidArgument, "missing client-id in metadata")
	}
	clientId := val[0]
	utils.Logger.Infof("client id is %s", clientId)

	s.streams[clientId] = stream

	// 接收客户端数据
	err := s.doRecv(stream)
	if err != nil {
		return err
	}
	return nil
}

// 接收grpc客户端数据
func (s *grpcServer) handleRecv() error {
	for _, server := range s.streams {
		stream := server
		go func() {
			s.doRecv(stream)
		}()
	}
	return nil
}

func (s *grpcServer) doRecv(stream v2.StreamDataService_AgentStreamServer) error {
	utils.Logger.Infof("doRecv %v", stream)
	for {
		recv, err := stream.Recv()
		utils.Logger.Infof("接收到客户端<%s>数据 %v", recv.ServerType, recv)
		// 错误处理
		if err != nil {
			log.Println(err)
			if grpc.Code(err) == codes.Canceled {
				stream = nil
			}

			if err == io.EOF {
				stream = nil
			}
			utils.Logger.Infof("recv error : %v", err)
		}
		if stream == nil {
			continue
		}
		// 客户端数据处理
		//_, isRequest := message.Message.(*v2.StreamMessage_Request)
		if _, ok := recv.Message.(*v2.StreamMessage_Response); ok {
			err := s.handleResponse(recv)
			if err != nil {
				utils.Logger.Infof("handleResponse error : %v", err)
			}
		} else {
			s.handleRequest(recv)
		}
	}
}

// 处理grpc客户端响应数据
func (s *grpcServer) handleResponse(recv *v2.StreamMessage) error {
	err := s.sendMessage(recv)
	if err != nil {
		return err
	}
	return nil
}

// 处理grpc客户端请求数据
func (s *grpcServer) handleRequest(recv *v2.StreamMessage) {
	err := s.sendMessage(recv)
	if err != nil {
		utils.Logger.Infof("handleRequest error : %v", err)
	}
}

// 处理转发请求
func (s *grpcServer) forwardRequest() {
	for {
		select {
		case request := <-s.streamData.Requests:
			utils.Logger.Info(request)
			if err := s.sendMessage(request); err != nil {
				utils.Logger.Info(err)
			}
		}
	}

}

// 发送消息到grpc客户端
func (s *grpcServer) sendMessage(message *v2.StreamMessage) error {

	stream := s.getServerStream(message.GetServerType())

	if stream == nil {
		return errors.New("stream does not exist")
	}
	utils.Logger.Infof("向客户端<%s>发送消息 %v", message.GetServerType(), message)
	err := stream.Send(message)

	if err != nil {
		return err
	}

	return nil
}

func (s *grpcServer) getServerStream(serverType string) v2.StreamDataService_AgentStreamServer {
	stream, ok := s.streams[serverType]
	if ok {
		return stream
	}

	return nil
}

// Request 接收转发请求
func (s *grpcServer) Request(message *v2.StreamMessage, requestId uint64) (*v2.StreamMessage, error) {
	_, isRequest := message.Message.(*v2.StreamMessage_Request)

	if isRequest {
		atomic.AddInt64(&reqId, 1)
		message.RequestId = uint64(reqId)
	} else {
		message.RequestId = requestId
	}

	s.streamData.Requests <- message
	if isRequest {
		for {
			select {
			case <-time.After(5 * time.Second):
				return nil, errors.New("request timeout")
			case msg := <-s.streamData.Responses:
				// TODO 这里的处理方式还需要优化
				utils.Logger.Infof("msg %v", msg)
				utils.Logger.Infof("message %v", message)
				if msg.RequestId == message.RequestId {
					return msg, nil
				}
			}
		}
	}
	return nil, nil
}

func (s *grpcServer) Start() error {
	go func() {
		s.forwardRequest()
	}()

	listener, _ := net.Listen("tcp", s.address)
	grpcserver := grpc.NewServer()
	// 注册服务方法
	v2.RegisterStreamDataServiceServer(grpcserver, s)

	if err := grpcserver.Serve(listener); err != nil {
		utils.Logger.Info(err)
		return err
	}

	return nil
}

func (s *grpcServer) Stop() {
	// TODO: 关闭连接
}
